{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "from utils.hand_model_lite import HandModelMJCFLite\n",
    "import numpy as np\n",
    "import transforms3d\n",
    "import torch\n",
    "import trimesh\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh_path = \"../data/meshdata\"\n",
    "data_path = \"../data/dataset\"\n",
    "\n",
    "use_visual_mesh = False\n",
    "\n",
    "hand_file = \"mjcf/shadow_hand_vis.xml\" if use_visual_mesh else \"mjcf/shadow_hand_wrist_free.xml\"\n",
    "\n",
    "joint_names = [\n",
    "    'robot0:FFJ3', 'robot0:FFJ2', 'robot0:FFJ1', 'robot0:FFJ0',\n",
    "    'robot0:MFJ3', 'robot0:MFJ2', 'robot0:MFJ1', 'robot0:MFJ0',\n",
    "    'robot0:RFJ3', 'robot0:RFJ2', 'robot0:RFJ1', 'robot0:RFJ0',\n",
    "    'robot0:LFJ4', 'robot0:LFJ3', 'robot0:LFJ2', 'robot0:LFJ1', 'robot0:LFJ0',\n",
    "    'robot0:THJ4', 'robot0:THJ3', 'robot0:THJ2', 'robot0:THJ1', 'robot0:THJ0'\n",
    "]\n",
    "translation_names = ['WRJTx', 'WRJTy', 'WRJTz']\n",
    "rot_names = ['WRJRx', 'WRJRy', 'WRJRz']\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hand_model = HandModelMJCFLite(\n",
    "    hand_file,\n",
    "    \"mjcf/meshes\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grasp_code_list = []\n",
    "for code in os.listdir(data_path):\n",
    "    grasp_code_list.append(code[:-4])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grasp_code = random.choice(grasp_code_list)\n",
    "grasp_data = np.load(\n",
    "    os.path.join(data_path, grasp_code+\".npy\"), allow_pickle=True)\n",
    "object_mesh_origin = trimesh.load(os.path.join(\n",
    "    mesh_path, grasp_code, \"coacd/decomposed.obj\"))\n",
    "print(grasp_code)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "index = random.randint(0, len(grasp_data) - 1)\n",
    "\n",
    "\n",
    "qpos = grasp_data[index]['qpos']\n",
    "rot = np.array(transforms3d.euler.euler2mat(\n",
    "    *[qpos[name] for name in rot_names]))\n",
    "rot = rot[:, :2].T.ravel().tolist()\n",
    "hand_pose = torch.tensor([qpos[name] for name in translation_names] + rot + [qpos[name]\n",
    "                         for name in joint_names], dtype=torch.float, device=\"cpu\").unsqueeze(0)\n",
    "hand_model.set_parameters(hand_pose)\n",
    "hand_mesh = hand_model.get_trimesh_data(0)\n",
    "object_mesh = object_mesh_origin.copy().apply_scale(grasp_data[index][\"scale\"])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(hand_mesh+object_mesh).show()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.16"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "2d4291a8a48b28165ac7232d5c96e1282ed596e61d078e48e9495becf16bfc7d"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
